import numpy as np
from scipy.interpolate import interp2
from scipy.special import gamma


def aimFcn_2(X, option, data):
    """
    基于B样条曲线的路径优化目标函数
    参数:
        X: 决策变量数组
        option: 算法参数配置
        data: 环境数据字典
    返回:
        total_cost: 总成本
        constraint_cost: 约束惩罚成本
    """
    X = X.flatten()  # 确保X为一维数组
    num_control_points = int((len(X) - data['E0'].shape[0]) / 3)
    required_elements = 3 * num_control_points
    
    if len(X) < required_elements:
        raise ValueError(f"输入向量X长度不足，需要至少{required_elements}个元素")
    
    # 提取控制点并构建路径点
    control_points = X[:required_elements].reshape(num_control_points, 3)
    path_points = np.vstack([data['S'], control_points, data['E0']])
    smoothed_path = cubic_bspline(path_points)
    
    # 初始化成本项
    distance_cost = 0.0
    height_cost = 0.0
    smooth_cost = 0.0
    constraint_cost = 0.0
    
    # 计算路径成本
    for j in range(smoothed_path.shape[0] - 1):
        # 飞行距离成本
        delta = smoothed_path[j+1] - smoothed_path[j]
        distance_cost += np.linalg.norm(delta)
        
        # 高度变化成本
        height_cost += abs(smoothed_path[j+1, 2] - smoothed_path[j, 2])
        
        # 路径平滑度成本
        if j > 0:
            vec1 = smoothed_path[j] - smoothed_path[j-1]
            vec2 = smoothed_path[j+1] - smoothed_path[j]
            cos_theta = np.dot(vec1[:2], vec2[:2]) / (np.linalg.norm(vec1[:2]) * np.linalg.norm(vec2[:2]) + 1e-8)
            cos_theta = np.clip(cos_theta, -1, 1)  # 防止数值误差
            angle = np.arccos(cos_theta)
            smooth_cost += angle
        
        # 碰撞检测
        current_point = smoothed_path[j]
        in_obstacle, _ = check_collision(current_point, data)
        if in_obstacle:
            constraint_cost += 20.0
        
        # 高度约束检查
        terrain_height = interp2(data['map_x'], data['map_y'], data['map_z'],
                                current_point[0], current_point[1], method='linear')
        if current_point[2] - terrain_height < data['minH']:
            constraint_cost += 20.0
        if current_point[2] > data['maxH']:
            constraint_cost += 20.0
    
    # 水平转弯角约束
    for j in range(1, smoothed_path.shape[0] - 1):
        vec1 = smoothed_path[j, :2] - smoothed_path[j-1, :2]
        vec2 = smoothed_path[j+1, :2] - smoothed_path[j, :2]
        cos_theta = np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2) + 1e-8)
        cos_theta = np.clip(cos_theta, -1, 1)
        hor_angle = np.arccos(cos_theta)
        if hor_angle > np.deg2rad(45):
            constraint_cost += 20.0
    
    # 爬升坡度约束
    for j in range(1, smoothed_path.shape[0]):
        hor_dist = np.linalg.norm(smoothed_path[j, :2] - smoothed_path[j-1, :2])
        ver_diff = abs(smoothed_path[j, 2] - smoothed_path[j-1, 2])
        if hor_dist < 1e-8:
            slope = 0
        else:
            slope = np.arctan(ver_diff / hor_dist)
        if slope > np.deg2rad(45):
            constraint_cost += 20.0
    
    # 加权总成本计算
    k1 = 1/20  # 距离权重
    k2 = 1/15  # 高度权重
    k3 = 1     # 平滑度权重
    total_cost = k1 * distance_cost + k2 * height_cost + k3 * smooth_cost + constraint_cost
    
    return total_cost, constraint_cost


def check_collision(point, data):
    """检查点是否与障碍物碰撞"""
    in_obstacle = False
    obstacle_height = 0.0
    x, y, z = point
    
    # 检查山地障碍物
    for i in range(data['Ob'].shape[0]):
        x0, y0, h, xs, ys = data['Ob'][i]
        gauss_val = h * np.exp(-((x - x0)/xs)**2 - ((y - y0)/ys)**2)
        if z <= gauss_val:
            in_obstacle = True
            obstacle_height = gauss_val
            return in_obstacle, obstacle_height
    
    # 检查城市障碍物
    if 'urban_obstacles' in data:
        for obs in data['urban_obstacles']:
            x_min, x_max, y_min, y_max, h = obs
            if x_min <= x <= x_max and y_min <= y <= y_max and z <= h:
                in_obstacle = True
                obstacle_height = h
                return in_obstacle, obstacle_height
    
    return in_obstacle, obstacle_height


def cubic_bspline(control_points):
    """三次B样条曲线平滑"""
    k = 3  # 三次B样条
    n = control_points.shape[0] - 1  # 控制点数量-1
    t = np.linspace(0, n + k + 1, 1000)  # 参数t
    
    # 均匀非周期性节点向量
    u = np.zeros(n + k + 2)
    for i in range(n + k + 2):
        if i <= k:
            u[i] = 0
        elif i > k and i <= n + 1:
            u[i] = i - k
        else:
            u[i] = n - k + 1
    
    # 计算B样条曲线点
    smoothed_path = np.zeros((len(t), 3))
    for j in range(len(t)):
        tj = t[j]
        for dim in range(3):
            val = 0.0
            for i in range(n + 1):
                basis = basis_function(i, k, tj, u)
                val += control_points[i, dim] * basis
            smoothed_path[j, dim] = val
    
    return smoothed_path


def basis_function(i, k, t, u):
    """Cox-deBoor递归公式计算B样条基函数"""
    if k == 1:
        return 1.0 if u[i] <= t < u[i+1] else 0.0
    
    den1 = u[i + k] - u[i]
    den2 = u[i + k + 1] - u[i + 1]
    
    term1 = 0.0
    if den1 > 1e-8:
        term1 = (t - u[i]) / den1 * basis_function(i, k-1, t, u)
    
    term2 = 0.0
    if den2 > 1e-8:
        term2 = (u[i + k + 1] - t) / den2 * basis_function(i+1, k-1, t, u)
    
    return term1 + term2